#!/usr/bin/env python
# -*- coding=utf-8 -*-
"""
@author: xingwg
@license: (C) Copyright 2020-2025.
@contact: xingweiguo@chinasvt.com
@project: boya-reid
@file: data_sampler.py
@time: 2020/9/13 0:53
@desc:
"""
import copy
import random
import glog
import numpy as np
from torch.utils.data.sampler import Sampler
from collections import defaultdict


class RandomIdentitySampler(Sampler):
    """
    随机采样N个ID，然后对每个ID随机采样K个图片，因此batch_size=N*K
    """
    def __init__(self, dataset, batch_size, num_instances):
        """
        构造函数
        :param dataset: list of (img_path, person_id, cam_id).
        :param batch_size: number of examples in a batch.
        :param num_instances: number of instances per identity in a batch.
        """
        super(RandomIdentitySampler, self).__init__(dataset)
        self.dataset = dataset
        self.batch_size = batch_size
        self.num_instances = num_instances
        self.num_pids_per_batch = self.batch_size // self.num_instances
        self.idx_dict = defaultdict(list)  # dict with list value

        for idx, (_, pid, _) in enumerate(self.dataset):
            self.idx_dict[pid].append(idx)
        self.pids = list(self.idx_dict.keys())

        # estimate number of examples in an epoch
        self.length = 0
        for pid in self.pids:
            idxs = self.idx_dict[pid]
            num = len(idxs)
            if num < self.num_instances:
                num = self.num_instances
            self.length += num - num % self.num_instances

    def __iter__(self):
        batch_idxs_dict = defaultdict(list)
        for pid in self.pids:
            idxs = copy.deepcopy(self.idx_dict[pid])
            if len(idxs) < self.num_instances:
                idxs = np.random.choice(idxs, size=self.num_instances, replace=True)
            random.shuffle(idxs)
            batch_idxs = []
            for idx in idxs:
                batch_idxs.append(idx)
                if len(batch_idxs) == self.num_instances:
                    batch_idxs_dict[pid].append(batch_idxs)
                    batch_idxs = []

        avai_pids = copy.deepcopy(self.pids)
        final_idxs = []

        while len(avai_pids) >= self.num_pids_per_batch:
            # 基于ID均匀采样
            selected_pids = random.sample(avai_pids, self.num_pids_per_batch)
            for pid in selected_pids:
                batch_idxs = batch_idxs_dict[pid].pop(0)
                final_idxs.extend(batch_idxs)
                if len(batch_idxs_dict[pid]) == 0:
                    avai_pids.remove(pid)

        return iter(final_idxs)

    def __len__(self):
        return self.length


class RandomIdentitySamplerV2(Sampler):
    """
    triplet 采样实现了过采样：对样本数<100的id进行过采样，
    让这些id的图片列表复制了两倍。这种做法在初赛的时候有1.5个点的提升
    https://github.com/anxiangsir/NAIC_reid_challenge_rank14_rank25/blob/master/data/samplers/triplet_sampler.py
    Randomly sample N identities, then for each identity,
    randomly sample K instances, therefore batch size is N*K.
    Args:
    - data_source (list): list of (img_path, pid, camid).
    - num_instances (int): number of instances per identity in a batch.
    - batch_size (int): number of examples in a batch.
    """

    def __init__(self, data_source, batch_size, num_instances):
        self.data_source = data_source
        self.batch_size = batch_size
        self.num_instances = num_instances
        self.num_pids_per_batch = self.batch_size // self.num_instances
        self.index_dic = defaultdict(list)
        for index, (_, pid, _) in enumerate(self.data_source):
            self.index_dic[pid].append(index)
        self.pids = list(self.index_dic.keys())

        # estimate number of examples in an epoch
        self.length = 0
        for pid in self.pids:
            idxs = self.index_dic[pid]
            num = len(idxs)
            if num < self.num_instances:
                num = self.num_instances
            self.length += num - num % self.num_instances

    def __iter__(self):
        batch_idxs_dict = defaultdict(list)

        pidslib = []
        for pid in self.pids:
            idxs = copy.deepcopy(self.index_dic[pid])

            # 对图片数量小于100的ID，图片数量扩大两倍，
            # 在基于图片数量采样的情况下，增加被采样概率
            pidslib += [pid] * len(idxs) * 2 if len(idxs) < 100 else [pid] * len(idxs)

            # 在采样数量不足num_instances时，随机复用数据
            if len(idxs) < self.num_instances:
                idxs = np.random.choice(idxs, size=self.num_instances, replace=True)
            # elif len(idxs) < 100:  # <100数据复制两倍
            #     idxs = np.random.choice(idxs, size=len(idxs) * 2, replace=True)
            # pidslib += [pid] * len(idxs)
            random.shuffle(idxs)
            batch_idxs = []
            for idx in idxs:
                batch_idxs.append(idx)
                if len(batch_idxs) == self.num_instances:
                    batch_idxs_dict[pid].append(batch_idxs)
                    batch_idxs = []

        avai_pids = copy.deepcopy(self.pids)
        final_idxs = []
        random.shuffle(pidslib)

        while len(avai_pids) >= self.num_pids_per_batch:
            # selected_pids = random.sample(avai_pids, self.num_pids_per_batch)
            selected_pidset = set()
            while len(selected_pidset) < self.num_pids_per_batch:
                selectedpid = pidslib.pop(0)
                if selectedpid in avai_pids:
                    selected_pidset.add(selectedpid)
            selected_pids = list(selected_pidset)

            for pid in selected_pids:
                batch_idxs = batch_idxs_dict[pid].pop(0)
                final_idxs.extend(batch_idxs)
                if len(batch_idxs_dict[pid]) == 0:
                    avai_pids.remove(pid)

        self.length = len(final_idxs)
        leftlst = []
        for pid in avai_pids:
            leftlst.append(len(batch_idxs_dict[pid]))
        glog.info("left number:{}".format(leftlst))
        return iter(final_idxs)

    def __len__(self):
        return self.length
